"""

This script takes in files that represent classifications of user-years into different clusters, and aggregates their data into:


	1. Count of users in each cluster by offset year and tenure, e.g.
		tenure	offset	c1	c2	c3
		0	0	131	21	30
		1	0	131	21	30
		1	1	100	10	3
		...
		9	0	131	21	30
		9	1	100	10	3
		...
		9	9	131	30	10

	2. Proportion of users in each cluster by offset year and tenure, e.g.
		tenure	offset	p1	p2	p3
		0	0	0.9	0.07	0.03
		1	0	0.8	0.1	0.1
		1	1	0.7	0.2	0.1
		...
		9	0	0.9	0.07	0.03
		9	1	0.8	0.1	0.1
		...
		9	9	0.9	0.07	0.03

Each block of statistics gets its own output file. If the CURRENT switch is set to true in this script, there will be four blocks of statistics; 
one for each separated by current users (i.e. users who have made an edit in the most recent year, also set in script) and users who have not made 
an edit in the last year.

Specify an input directory to read input files from.
Specify an output directory to write output files to.
Specify the number of clusters in the model. This probably isn't necessary but is a sanity check in case no users are classified to a cluster.

There are lots of parameters for which elements to look for in this script that are not passed as arguments. The defaults here should be good for 
our project.

This script is well commented enough that reading through the code should answer any questions about how these statistics are compiled.

"""

import sys
import os
import datetime

def main():

	inputDir = sys.argv[1]
	outputDir = sys.argv[2]
	numclusters = int(sys.argv[3])


	# Make a list of input files and open them.

	inputFileList = os.listdir(inputDir)
	inputFileHandles = []
	
	for inputFileName in inputFileList:
		
		inputFilePath = inputDir+'/'+inputFileName
		
		if os.path.isfile(inputFilePath):
			inputFileHandles.append(open(inputFilePath, 'r'))


	# Initialize count and proportion by tenure and offset

	tenure_offset = {}		# Address tenure_offset[tenure][offset][cluster]
	proportion_tenure_offset = {}	# Address proportion_tenure_offset[tenure][offset][cluster]

	tenure_offset_current = {}		# The count for current users, if we are separating, addressed the same as above
	proportion_tenure_offset_current = {}	# The proportion for current users, if we are separating, addressed the same as above

	start_tenure = 0		# The first tenure we look at in the data, inclusive.
	end_tenure = 9			# The last tenure we look at in the data, inclusive. The largest tenure in our data in 2012 is 9
	
	# The stats compiled with these fields are only populated up to offsets that make sense (i.e. tenure 2 group has no offset 3)

	current = True			# If we separate the statistics by current and non-current users (True = separate)
	current_year = 2011		# The year (and later) that we consider the current year - we choose 2011 for this since halfway through 2012.

	for tenure_index in range(start_tenure, end_tenure + 1):
		
		tenure_offset[tenure_index] = {}
		proportion_tenure_offset[tenure_index] = {}

		if current:
			tenure_offset_current[tenure_index] = {}
			proportion_tenure_offset_current[tenure_index] = {}

		for offset_index in range(0, tenure_index + 1):

			tenure_offset[tenure_index][offset_index] = {}
			proportion_tenure_offset[tenure_index][offset_index] = {}

			if current:
				tenure_offset_current[tenure_index][offset_index] = {}
				proportion_tenure_offset_current[tenure_index][offset_index] = {}

			for cluster_index in range(1, numclusters + 1):

				tenure_offset[tenure_index][offset_index][cluster_index] = 0
				proportion_tenure_offset[tenure_index][offset_index][cluster_index] = None

				if current:
					tenure_offset_current[tenure_index][offset_index][cluster_index] = 0
					proportion_tenure_offset_current[tenure_index][offset_index][cluster_index] = None

	# Compile counts

	print timestamp() + " Processing input files in " + inputDir

	total_entries = 0		# The total number of entries we look at

	for inputFileHandle in inputFileHandles:

		print timestamp() + " Processing input file " + inputFileHandle.name

		# This keeps track of the data for a single user so we can compute tenure and cohort for entries
		# As a result those counters are a line behind the ones that don't use this field
		# They also need a call after the loop body to fill in the data for the last user, since it breaks on 
		# only changes in username.
		userdata = []

		for line in inputFileHandle:

			fields = parseLine(line)
				

			# Update count by tenure and offset

			if userdata and fields['user'] != userdata[0]['user']:
				updateTenureOffset(tenure_offset, tenure_offset_current, userdata, current_year, current)


			# Reset user data if new user and update with this line if not
			
			if userdata and fields['user'] != userdata[0]['user']:
				userdata = [fields]
			else:
				userdata.append(fields)

			
			# Increment the total entry counter

			total_entries += 1


		# Update information for the last user in the list

		updateTenureOffset(tenure_offset, tenure_offset_current, userdata, current_year, current)


		print timestamp() + " Finished processing input file " + inputFileHandle.name


	print timestamp() + " Finished processing input files in " + inputDir


	# Compute proportions

	print timestamp() + " Compiling proportions"

	# Compute tenure and offset proportion

	tenure_offset_total = 0		# Sanity check to make sure we counted all entries
	

	for tenure_index in tenure_offset:

		for offset_index in tenure_offset[tenure_index]:

			single_tenure_offset_total = 0		# The total for one tenure group's offset year we compute proportion with
			current_single_tenure_offset_total = 0	# Same if we are separating by current users

			# Figure out total users in this tenure group and offset year

			for cluster_index in tenure_offset[tenure_index][offset_index]:

				tenure_offset_total += tenure_offset[tenure_index][offset_index][cluster_index]
				single_tenure_offset_total += tenure_offset[tenure_index][offset_index][cluster_index]

				if current:
					tenure_offset_total += tenure_offset_current[tenure_index][offset_index][cluster_index]
					current_single_tenure_offset_total += tenure_offset_current[tenure_index][offset_index][cluster_index]

			# Compute the proportion based on total users in this tenure group and offset year

			for cluster_index in tenure_offset[tenure_index][offset_index]:

				if single_tenure_offset_total != 0:
					proportion_tenure_offset[tenure_index][offset_index][cluster_index] = float(tenure_offset[tenure_index][offset_index][cluster_index]) / float(single_tenure_offset_total)
				else:
					proportion_tenure_offset[tenure_index][offset_index][cluster_index] = 0

				if current:
					if current_single_tenure_offset_total != 0:
						proportion_tenure_offset_current[tenure_index][offset_index][cluster_index] = float(tenure_offset_current[tenure_index][offset_index][cluster_index]) / float(current_single_tenure_offset_total)
					else:
						proportion_tenure_offset_current[tenure_index][offset_index][cluster_index] = 0

	assert(tenure_offset_total == total_entries), "Tenure and offset is missing " + str(total_entries - tenure_offset_total) + " entries"


	# Write output files

	print timestamp() + " Writing output files"


	# Write tenure and offset count statistics file
	
	tenure_offset_filename = 'tenure_offset.stat'				# Filename to write offset and tenure count statistics to
	current_tenure_offset_filename = 'tenure_offset_current.stat'		# Filename to write same to for current users if separating
	tenure_offset_filepath = outputDir + '/' + tenure_offset_filename
	current_tenure_offset_filepath = outputDir + '/' + current_tenure_offset_filename

	print timestamp() + " Writing tenure and offset count statistics file to " + tenure_offset_filepath

	if current:
		print timestamp() + " Writing current tenure and offset count statistics file to " + current_tenure_offset_filepath

	tenure_offset_filehandle = open(tenure_offset_filepath, 'w')

	if current:
		current_tenure_offset_filehandle = open(current_tenure_offset_filepath, 'w')
	
	for tenure_index in sorted(tenure_offset):
		

		for offset_index in sorted(tenure_offset[tenure_index]):

			output = [str(tenure_index), str(offset_index)]

			if current:
				current_output = [str(tenure_index), str(offset_index)]

			for cluster_index in sorted(tenure_offset[tenure_index][offset_index]):

				output.append(str(tenure_offset[tenure_index][offset_index][cluster_index]))

				if current:
					current_output.append(str(tenure_offset_current[tenure_index][offset_index][cluster_index]))

			tenure_offset_filehandle.write('\t'.join(output) + '\n')

			if current:
				current_tenure_offset_filehandle.write('\t'.join(current_output) + '\n')
				

	tenure_offset_filehandle.close()
	
	if current:
		current_tenure_offset_filehandle.close()

	print timestamp() + " Finished writing tenure and offset count statistics file to " + tenure_offset_filepath

	if current:
		print timestamp() + " Finished writing current tenure and offset count statistics file to " + current_tenure_offset_filepath

	
	# Write tenure and offset proportion statistics file

	proportion_tenure_offset_filename = 'tenure_offset_proportion.stat'	# Filename to write offset and tenure proportion statistics to
	current_proportion_tenure_offset_filename = 'tenure_offset_proportion_current.stat' # Filename for same for current users if separating
	proportion_tenure_offset_filepath = outputDir + '/' + proportion_tenure_offset_filename
	current_proportion_tenure_offset_filepath = outputDir + '/' + current_proportion_tenure_offset_filename

	print timestamp() + " Writing tenure and offset proportion statistics file to " + proportion_tenure_offset_filepath

	if current:
		print timestamp() + " Writing current tenure and offset proportion statistics file to " + current_proportion_tenure_offset_filepath

	proportion_tenure_offset_filehandle = open(proportion_tenure_offset_filepath, 'w')

	if current:
		current_proportion_tenure_offset_filehandle = open(current_proportion_tenure_offset_filepath, 'w')

	for tenure_index in sorted(proportion_tenure_offset):
		

		for offset_index in sorted(proportion_tenure_offset[tenure_index]):

			output = [str(tenure_index), str(offset_index)]

			if current:
				current_output = [str(tenure_index), str(offset_index)]

			for cluster_index in sorted(proportion_tenure_offset[tenure_index][offset_index]):

				output.append(str(proportion_tenure_offset[tenure_index][offset_index][cluster_index]))

				if current:
					current_output.append(str(proportion_tenure_offset_current[tenure_index][offset_index][cluster_index]))

			proportion_tenure_offset_filehandle.write('\t'.join(output) + '\n')

			if current:
				current_proportion_tenure_offset_filehandle.write('\t'.join(current_output) + '\n')

	proportion_tenure_offset_filehandle.close()

	if current:
		current_proportion_tenure_offset_filehandle.close()

	print timestamp() + " Finished writing tenure and offset proportion statistics file to " + proportion_tenure_offset_filepath

	if current:
		print timestamp() + " Finished current writing tenure and offset proportion statistics file to " + current_proportion_tenure_offset_filepath


def updateTenureOffset(tenure_offset, current_tenure_offset, userdata, currentyear, current):
	startyear = userdata[0]['year']
	endyear = userdata[len(userdata)-1]['year']
	usertenure = endyear - startyear

	if current and endyear >= currentyear:

		for entry in userdata:
			current_tenure_offset[usertenure][entry['offset']][entry['cluster']] += 1

	else:
		for entry in userdata:
			tenure_offset[usertenure][entry['offset']][entry['cluster']] += 1

def parseLine(line):
	
	tokens = line.strip().split('\t')

	fields = {}
	fields['user'] = tokens[0]
	fields['offset'] = int(tokens[1])
	fields['year'] = int(tokens[2])
	fields['cluster'] = int(tokens[3])

	return fields

def timestamp():
	
	return datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')

main()
